#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from PyQt5 import Qt
from gnuradio import qtgui
from gnuradio.filter import firdes
import sip
from gnuradio import gr
from gnuradio import blocks
import sys
import signal
from gnuradio import uhd
import time
import math
import numpy as np


class observe_wifi_spectrum(gr.top_block, Qt.QWidget):

    def __init__(self):
        gr.top_block.__init__(self, "WiFi Spectrum Observer")
        Qt.QWidget.__init__(self)
        self.setWindowTitle("WiFi Spectrum Observer")

        # Variables
        self.samp_rate = 3e5
        self.center_freq = 200e6
        self.gain = 0  # USRP gain

        # USRP Source
        self.uhd_usrp_source_0 = uhd.usrp_source(
            ",".join(("=serial=31A92E2", "")),
            uhd.stream_args(
                cpu_format="fc32",
                args='',
                channels=list(range(0, 1)),
            ),
        )
        self.uhd_usrp_source_0.set_center_freq(self.center_freq, 0)
        self.uhd_usrp_source_0.set_gain(self.gain, 0)
        self.uhd_usrp_source_0.set_antenna('TX/RX', 0)
        self.uhd_usrp_source_0.set_samp_rate(self.samp_rate)
        self.uhd_usrp_source_0.set_auto_dc_offset(True)
        self.uhd_usrp_source_0.set_auto_iq_balance(True)

        # Frequency Sink
        self.qtgui_freq_sink_x_0 = qtgui.freq_sink_c(
            2048,  # FFT size
            firdes.WIN_BLACKMAN_hARRIS,
            self.center_freq,
            self.samp_rate,
            "",
            1
        )
        self.qtgui_freq_sink_x_0.set_update_time(1)
        self.qtgui_freq_sink_x_0.set_y_axis(-100, 0)
        self.qtgui_freq_sink_x_0.set_y_label('Relative Gain', 'dB')
        self.qtgui_freq_sink_x_0.enable_grid(True)
        self.qtgui_freq_sink_x_0.enable_autoscale(True)

        # Power Measurement using magnitude squared and moving average
        self.mag_squared = blocks.complex_to_mag_squared(1)
        self.moving_avg = blocks.moving_average_ff(100, 1 / 100, 4000)

        # Probe for power measurement
        self.probe = blocks.probe_signal_f()

        # Connections
        self.connect((self.uhd_usrp_source_0, 0), (self.qtgui_freq_sink_x_0, 0))
        self.connect((self.uhd_usrp_source_0, 0), (self.mag_squared, 0))
        self.connect((self.mag_squared, 0), (self.moving_avg, 0))
        self.connect((self.moving_avg, 0), (self.probe, 0))

        # Qt GUI
        self._qtgui_freq_sink_x_0_win = sip.wrapinstance(self.qtgui_freq_sink_x_0.pyqwidget(), Qt.QWidget)
        self.setLayout(Qt.QVBoxLayout())
        self.layout().addWidget(self._qtgui_freq_sink_x_0_win)

        # Setup timer to print power
        self.timer = Qt.QTimer()
        self.timer.timeout.connect(self.print_power)
        self.timer.start(1000)  # Update every 1000ms

    def print_power(self):
        try:
            power = self.probe.level()
            if power > 0:
                power_dbm = 10 * math.log10(power) + 30  # Convert to dBm
                print(f"Received Power: {power_dbm:.2f} dBm")
            else:
                print("Received Power: -inf dBm")
        except Exception as e:
            print(f"Error reading power: {str(e)}")

    def closeEvent(self, event):
        self.timer.stop()
        event.accept()


def main():
    qapp = Qt.QApplication(sys.argv)
    tb = observe_wifi_spectrum()
    tb.start()
    tb.show()

    def sig_handler(sig=None, frame=None):
        Qt.QApplication.quit()

    signal.signal(signal.SIGINT, sig_handler)
    signal.signal(signal.SIGTERM, sig_handler)

    qapp.exec_()


if __name__ == '__main__':
    main()